package com.wondersgroup.aida.ragdemo;

import com.wondersgroup.aida.tools.ModelUtil;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentParser;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.TextDocumentParser;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.openai.OpenAiEmbeddingModel;
import dev.langchain4j.store.embedding.redis.RedisEmbeddingStore;

import java.net.URISyntaxException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 * @author: 紫麒麟vip
 * @create: 2025/2/16
 * Description:
 */
public class MeituanRagLoader {
    public static void main(String[] args) throws URISyntaxException {
        // 1. 读取本地知识文件
        Path documentPaht = Paths.get(MeituanRagLoader.class.getClassLoader().getResource("meituan-questions.txt").toURI());
        DocumentParser documentParser = new TextDocumentParser();
        Document document = FileSystemDocumentLoader.loadDocument(documentPaht, documentParser);
        // 2. 把知识文件分解成一个一个的知识条目
        DocumentSplitter splitter = new MyDocumentSplitter();
        List<TextSegment> segments = splitter.split(document);
        // 3. 对每个条目进行文本向量化，并且保存到Redis中。
        EmbeddingModel embeddingModel = OpenAiEmbeddingModel.builder()
                .apiKey(ModelUtil.API_KEY_0PENAI)
                .baseUrl(ModelUtil.BASE_URI_OPENAI)
                .build();
        RedisEmbeddingStore embeddingStore = RedisEmbeddingStore.builder()
                .host("127.0.0.1")
                .port(6379)
                .dimension(1536)
                .indexName("meituan-rag")
                .build();

        List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

        embeddingStore.addAll(embeddings,segments);
    }
}
